{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2023 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long PDF Q&A with Gemini 2.0\n",
    "\n",
    "The goal of this notebook is to extract specific information from a large PDF by using Gemini 2.0.\n",
    "\n",
    "In this notebook, you will:\n",
    " - Use Gemini to answer a specific question contained in a PDF document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import python packages\n",
    "\n",
    "from typing import Iterable\n",
    "import io\n",
    "import time\n",
    "\n",
    "import vertexai\n",
    "from vertexai.preview.generative_models import (\n",
    "    GenerationResponse,\n",
    "    GenerativeModel,\n",
    "    HarmBlockThreshold,\n",
    "    HarmCategory,\n",
    "    Part\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Include information about your project in the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = \"[your-project-id]\"  # Replace with your project ID\n",
    "LOCATION = \"us-central1\"  # Replace with your location\n",
    "MODEL_NAME = \"gemini-2.0-flash-001\"  # Replace with model name\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
    "model = GenerativeModel(MODEL_NAME)\n",
    "BLOCK_LEVEL = HarmBlockThreshold.BLOCK_ONLY_HIGH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "Use the document above to answer the question below. Follow the Instructions and Suggestions below as a guide to answering the question.\n",
    "<Instructions>\n",
    "- First, analyze the question below and return which variables need to be analyzed, from what time period (example: second quarter of 2020), and any other details present in the question.\n",
    "- Then return an analysis of what is asked in the question.\n",
    "- Finally, carefully analyze the document above and answer the question below completely and correctly, using the variables determined in the previous step.\n",
    "- Explain how you arrived at this result.\n",
    "- Answer ONLY what was asked.\n",
    "<Instructions>\n",
    "<Suggestions>\n",
    "- The document above is a financial report with various tables, graphs, infographics, lists, and additional information in text.\n",
    "- PAY VERY CLOSE ATTENTION to the legends of the graphs and the COLORS of the graphs to answer the question below. The colors may indicate which information is important to answer the question.\n",
    "- The color of the graph legends represents the color of the graph bars.\n",
    "- Use ONLY this document as context to answer the question below.\n",
    "</Suggestions>\n",
    "<Question>\n",
    "{question}\n",
    "</Question>\n",
    "answer:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(\n",
    "    prompt: list,\n",
    "    max_output_tokens: int = 2048,\n",
    "    temperature: int = 2,\n",
    "    top_p: float = 0.4,\n",
    "    stream: bool = False,\n",
    ") -> GenerationResponse | Iterable[GenerationResponse]:\n",
    "    \"\"\"\n",
    "    Function to generate response using Gemini 2.0\n",
    "\n",
    "    Args:\n",
    "        prompt:\n",
    "            List of prompt parts\n",
    "        max_output_tokens:\n",
    "            Max Output tokens\n",
    "        temperature:\n",
    "            Temperature for the model\n",
    "        top_p:\n",
    "            Top-p for the model\n",
    "        stream:\n",
    "            Strem results?\n",
    "\n",
    "    Returns:\n",
    "        Model response\n",
    "\n",
    "    \"\"\"\n",
    "    responses = model.generate_content(\n",
    "        prompt,\n",
    "        generation_config={\n",
    "            \"max_output_tokens\": max_output_tokens,\n",
    "            \"temperature\": temperature,\n",
    "            \"top_p\": top_p,\n",
    "        },\n",
    "        safety_settings={\n",
    "            HarmCategory.HARM_CATEGORY_HATE_SPEECH: BLOCK_LEVEL,\n",
    "            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: BLOCK_LEVEL,\n",
    "            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: BLOCK_LEVEL,\n",
    "            HarmCategory.HARM_CATEGORY_HARASSMENT: BLOCK_LEVEL,\n",
    "        },\n",
    "        stream=stream,\n",
    "    )\n",
    "\n",
    "    return responses\n",
    "\n",
    "\n",
    "def retry_generate(pdf_document: Part, prompt: str, question: str):\n",
    "    predicted = False\n",
    "    while not predicted:\n",
    "        try:\n",
    "            response = generate(\n",
    "                prompt=[pdf_document, prompt.format(question=question)]\n",
    "            )\n",
    "        except Exception as e:\n",
    "            print(\"sleeping for 2 seconds ...\")\n",
    "            print(e)\n",
    "            time.sleep(2)\n",
    "        else:\n",
    "            predicted = True\n",
    "\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample questions\n",
    "\n",
    "In the next cell, include information about your question and the pdf_path.  \n",
    "\n",
    "**(Optional)**  \n",
    "If you are using Colab to test this notebook, you can try the following code to upload your PDF files.  \n",
    "```python\n",
    "from google.colab import files\n",
    "files.upload()\n",
    "```\n",
    "\n",
    "You can uncomment the code in the cell to use this method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from google.colab import files\n",
    "# files.upload()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"From the Consolidated Balance Sheet, what was the difference between the total assets from 2022 to 2023?\"\n",
    "pdf_path = \"./Cymbal Bank - Financial Statements.pdf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "## Analysis of the Question:\n",
      "\n",
      "The question asks for the difference in **total assets** between the years **2022** and **2023** from the **Consolidated Balance Sheet**. This requires locating the relevant section within the document and identifying the values associated with each year. \n",
      "\n",
      "\n",
      "## Locating the Information:\n",
      "\n",
      "1. **Consolidated Balance Sheet:** The document provides a \"Consolidated Balance Sheet\" table which contains financial data for the years 2022 and 2023.\n",
      "2. **Total Assets:** We need to identify the row labeled \"Total assets\" within the table. \n",
      "3. **Values for 2022 and 2023:**  We will find the corresponding values under the \"12/31/2022\" and \"12/31/2023\" columns.\n",
      "\n",
      "\n",
      "## Calculation:\n",
      "\n",
      "1. **2023 Total Assets:**  $2,238,274 million \n",
      "2. **2022 Total Assets:** $2,281,868 million\n",
      "3. **Difference:** $2,281,868 million - $2,238,274 million = $43,594 million\n",
      "\n",
      "\n",
      "## Answer:\n",
      "\n",
      "The difference in total assets between 2022 and 2023 according to the Consolidated Balance Sheet is **$43,594 million**. This indicates a decrease in total assets from 2022 to 2023. \n",
      "\n"
     ]
    }
   ],
   "source": [
    "with open(pdf_path, \"rb\") as fp:\n",
    "    pdf_document = Part.from_data(data=fp.read(), mime_type=\"application/pdf\")\n",
    "\n",
    "response = retry_generate(pdf_document, prompt, question)\n",
    "print(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "workbench-notebooks.m128",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m128"
  },
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
